# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/17 10:52 
@Author  : ZhangShenao 
@File    : 3.StructuredToolChat结构化工具对话代理.py 
@Desc    : 使用Structured Tool Chat结构化工具对话代理
"""
import asyncio
import os

import dotenv
from langchain.agents import initialize_agent, AgentType
from langchain_community.agent_toolkits import PlayWrightBrowserToolkit
from langchain_community.tools.playwright.utils import create_async_playwright_browser
from langchain_openai import ChatOpenAI

# 加载环境变量
dotenv.load_dotenv()

# 创建PlayWrightBrowserToolkit工具箱
async_browser = create_async_playwright_browser()
toolkit = PlayWrightBrowserToolkit.from_browser(async_browser=async_browser)
tools = toolkit.get_tools()

# 创建LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))

# 初始化Agent
# 使用PlayWrightBrowserToolkit工具箱,采用STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION代理模式
agent_chain = initialize_agent(
    tools=tools,  # 指定需要使用到的工具箱
    llm=llm,  # 指定LLM
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,  # 指定代理模式为Structured Tool Chat结构化工具对话
    verbose=True,  # 输出调试信息
)


async def run():
    """异步执行函数"""

    # 异步执行Agent,获取结果
    response = await agent_chain.arun('What are the headers on https://www.imooc.com/?')
    print(response)


# 执行异步函数
asyncio.run(run())
